from Crypto.Cipher import AES
from Crypto.Random import get_random_bytes
import base64
import struct

# 常量定义（与原JS完全一致）
jf = 16  # 随机值n的长度（16字节）
L_ = 16  # 随机值r的长度（16字节，作为AES-IV）
Fde = 15536  # 原JS未使用，保留定义
Lde = 256    # 原JS未使用，保留定义
Nde = {      # 原JS未使用，保留定义
    "words": [2037738088, 2036823142, 1919247475, 2003007078],
    "sigBytes": 16
}


def kde(e: bytes) -> str:
    return ''.join(chr(byte) for byte in e)


def dinit(l: list[int]) -> dict:
    return {
        "words": l,
        "sigBytes": 4 * len(l)  # 每个word占4字节
    }


def parse(hex_str: str) -> dict:
    str_len = len(hex_str)
    words = []
    x = 0
    while x < str_len:
        byte_val = int(hex_str[x:x+2], 16)
        word_idx = x // 8
        shift = 24 - (x % 8) * 4  # 第0个字节占24-31位，第1个占20-23位...
        
        if len(words) <= word_idx:
            words.append(0)
        
        words[word_idx] |= byte_val << shift
        x += 2  # 处理下一个字节
    
    actual_sig_bytes = str_len // 2
    return {
        "words": words,
        "sigBytes": actual_sig_bytes
    }


def wordarray_to_bytes(wordarray: dict) -> bytes:
    words = wordarray["words"]
    sig_bytes = wordarray["sigBytes"]
    byte_list = []
    
    for word in words:
        unsigned_word = word & 0xFFFFFFFF  # 处理负数（转为32位无符号）
        word_bytes = struct.pack('>I', unsigned_word)
        byte_list.extend(word_bytes)
    
    return bytes(byte_list[:sig_bytes])

#加密
def AESencrypt(plain_text: str) -> str:
    try:
        n = get_random_bytes(jf)
        r = get_random_bytes(L_)
        
        key_wordarray = {
            "words": [1141956472, 1820211443, 657019859, 565121240, 
                      76437628, -1107606939, -1647905496, -1693725629],
            "sigBytes": 32  # 8个word × 4字节 = 32字节
        }
        key = wordarray_to_bytes(key_wordarray)
        
        plain_bytes = plain_text.encode('utf-8')
        plain_hex = plain_bytes.hex()  # 对应原JS的a变量
        plain_wordarray = parse(plain_hex)  # 对应原JS的s变量
        plain_data = wordarray_to_bytes(plain_wordarray)
        
        r_hex = r.hex()  # 对应原JS的l变量
        iv_wordarray = parse(r_hex)  # 对应原JS的IV变量
        iv = wordarray_to_bytes(iv_wordarray)
        
        cipher = AES.new(key, AES.MODE_CBC, iv)
        pad_len = AES.block_size - (len(plain_data) % AES.block_size)
        if pad_len == 0:
            pad_len = AES.block_size  # 刚好整除时补一整块
        plain_data_padded = plain_data + bytes([pad_len]) * pad_len
        ciphertext = cipher.encrypt(plain_data_padded)
        
        combined = n + r + ciphertext
        return base64.b64encode(combined).decode('utf-8')
    
    except Exception as e:
        print(f"加密失败: {str(e)}")
        return plain_text

#解密
def AESdeencrypt(cipher_base64: str) -> str:
    try:
        combined = base64.b64decode(cipher_base64)
        
        r = combined[16:16+16]  # 对应原JS的r变量（IV来源）
        ciphertext_raw = combined[16+16:]  # 对应原JS的o变量（密文原始字节）
        
        key_wordarray = {
            "words": [1141956472, 1820211443, 657019859, 565121240, 
                      76437628, -1107606939, -1647905496, -1693725629],
            "sigBytes": 32
        }
        key = wordarray_to_bytes(key_wordarray)
        
        r_hex = r.hex()  # 对应原JS的i变量
        iv_wordarray = parse(r_hex)  # 对应原JS的s变量
        iv = wordarray_to_bytes(iv_wordarray)
        
        ciphertext_hex = ciphertext_raw.hex()  # 对应原JS的a变量
        ciphertext_wordarray = parse(ciphertext_hex)  # 对应原JS的l变量
        ciphertext = wordarray_to_bytes(ciphertext_wordarray)
        
        cipher = AES.new(key, AES.MODE_CBC, iv)
        decrypted_padded = cipher.decrypt(ciphertext)
        
        pad_len = decrypted_padded[-1]
        decrypted = decrypted_padded[:-pad_len]
        
        return decrypted.decode('utf-8')
    
    except Exception as e:
        print(f"解密失败: {str(e)}")
        return cipher_base64

